import os, sys
import re
import codecs
import numpy as np
import soundfile
import wave
import queue
from multiprocessing import cpu_count
import threading


def read_wav_scp(wav_scp):
    if not os.path.exists(wav_scp):
        return None
    id2wavdir = dict()
    with codecs.open(filename=wav_scp, mode='r') as f:
        line = f.readline().strip()
        while line is not None and line != "":
            key, dir = re.split(r"\s+", line, maxsplit=1)
            id2wavdir[key] = dir.strip()
            line = f.readline().strip()

    return id2wavdir


def read_segment(wav_dir):
    if not os.path.exists(wav_dir):
        print('wav {} is not exist !\n'.format(wav_dir))
        return None, None

    a_wave_fp = wave.open(wav_dir, "rb")
    params = a_wave_fp.getparams()
    nchannels, sampwidth, sample_rate, nframes = params[:4]
    str_data = a_wave_fp.readframes(nframes)
    a_wave_fp.close()
    wave_data = np.fromstring(str_data, dtype=np.short)
    return wave_data, nframes


def write_segment(wave_data, nframes, filename):
    sub_wave_fp = wave.open(filename, "wb")
    # 配置声道数、量化位数和取样频率
    sub_wave_fp.setnchannels(1)
    sub_wave_fp.setsampwidth(2)
    sub_wave_fp.setframerate(16000)
    # 将wav_data转换为二进制数据写入文件
    sub_wave_fp.writeframes(wave_data.tobytes())
    sub_wave_fp.close()


def splice_short_segment(
        dest_dir,
        queue,
        utt2dur,
        max_duration=55,
        thread_id=0):
    temp_keys, temp_durations = [], []
    temp_acc_dur = 0
    section_index = 0
    if len(utt2dur) > 0:
        for one_dur in utt2dur:
            key, duration = re.split(r"\s+", one_dur, maxsplit=1)
            temp_keys.append(key)
            temp_durations.append(float(duration))
            temp_acc_dur += float(duration)
            if (temp_acc_dur + len(temp_keys)*3) >= max_duration:
                total_interval = max_duration - np.sum(temp_durations)
                avg_interval = total_interval/len(temp_durations)
                interval_point_num = int(avg_interval*16000)
                section_wav = np.zeros(interval_point_num, dtype=np.short)
                infos = []
                for i in range(len(temp_keys)):
                    wave_data, nframes = read_segment(id2wavdirs[temp_keys[i]])
                    if wave_data is not None:
                        start_time = section_wav.shape[0]/16000
                        section_wav = np.hstack((section_wav, wave_data))
                        end_time = section_wav.shape[0]/16000
                        infos.append('[{}:{}:{}]'.format(temp_keys[i], start_time, end_time))
                        print('{}:{}\n'.format(start_time, end_time))
                        if i < len(temp_keys):
                            section_wav = np.hstack((section_wav, np.zeros(interval_point_num, dtype=np.short)))

                section_filename = str(thread_id).zfill(2) + "_" + str(section_index).zfill(10)
                write_segment(section_wav, -1, dest_dir + "//" + section_filename + ".wav")
                print('splice {} segment to one {}\n'.format(len(temp_keys), section_filename))
                queue.put('{} {}\n'.format(section_filename, " ".join(infos)))
                temp_acc_dur = 0
                temp_keys, temp_durations = [], []
                section_index += 1
                temp_acc_dur = 0


if __name__ == '__main__':
    wav_scp, utt2dur = "100.wav.scp", "100.utt2dur"
    dest_dir = "splice_short_wavs"
    if not os.path.exists(dest_dir):
        os.mkdir(dest_dir)
    id2wavdirs = read_wav_scp(wav_scp)
    info_queue = queue.Queue()
    num_cpu = 2
    durs = [line.strip() for line in codecs.open(filename=utt2dur, mode='r').readlines()]
    total_entry_num = len(durs)
    segment2section = wav_scp + "_map_table.txt"
    map_fp = codecs.open(filename=segment2section, mode='w')
    task_num_one_thread = int(total_entry_num / num_cpu) + 1
    t_objects = list()
    for thread_index in range(0, num_cpu - 1):
        entry_start = thread_index * task_num_one_thread
        entry_end = (thread_index + 1) * task_num_one_thread - 1
        if entry_end >= total_entry_num:
            entry_end = total_entry_num

        t = threading.Thread(target=splice_short_segment,
                                           args=(dest_dir, info_queue, durs[entry_start:entry_end], 55, thread_index))
        t_objects.append(t)

    for one_thread in t_objects:
        one_thread.start()

    for index_t in t_objects:
        index_t.join()

    while not info_queue.empty():
        one_info = info_queue.get()
        map_fp.write('{}'.format(one_info))
    map_fp.close()















